///////////////////////////////////////////////////////////////////////
// File:        EventManager.cpp
// Description: 
// Author:      Philip Kommnick
// Created:     Thu September 10 14:40:08 CEST 2009
//
// This program is open source. For license terms, see the LICENSE file.
//
///////////////////////////////////////////////////////////////////////

#include "stdafx.h"
#include "EventManager.h"
#include "TwainScanner.h"
#include "ImageFileSink.h"
#include "OcrSink.h"
#include "PrinterSink.h"
#include "resource.h"

// IUnknown
ULONG __stdcall EventManager::AddRef(void)
{
    return ++refCnt;
}

ULONG __stdcall EventManager::Release(void)
{
    return --refCnt;
}

HRESULT __stdcall EventManager::QueryInterface(REFIID riid, void **ppvObject)
{
    HRESULT hrRet = S_OK;
    *ppvObject = NULL;

    if (IsEqualIID(riid, IID_IUnknown)) {
        *ppvObject = (IUnknown*)this;
    } else if (IsEqualIID(riid, IID_IWiaEventCallback)) {
        *ppvObject = (IWiaEventCallback*)this;
    } else {
        hrRet = E_NOINTERFACE;
    }

    // If we're returning an interface pointer, AddRef() it.
    if ( S_OK == hrRet ) {
        ((IUnknown*) *ppvObject)->AddRef();
    }
    return hrRet;
}

////override less for GUIDS
//template<> bool std::less<GUID>::operator()(const GUID& _Left, const GUID& _Right) const
//{
//	if (_Left.Data1 != _Right.Data1)
//		return _Left.Data1 < _Right.Data1;
//	if (_Left.Data2 != _Right.Data2)
//		return _Left.Data2 < _Right.Data2;
//	if (_Left.Data3 != _Right.Data3)
//		return _Left.Data3 < _Right.Data3;
//	return reinterpret_cast<UINT32>(_Left.Data4) < reinterpret_cast<UINT32>(_Right.Data4);
//}

// IWiaEventCallback
HRESULT __stdcall EventManager::ImageEventCallback(const GUID *pEventGUID, BSTR bstrEventDescription, BSTR bstrDeviceID,
        BSTR bstrDeviceDescription, DWORD dwDeviceType, BSTR bstrFullItemName, ULONG *pulEventType, ULONG ulReserved)
{
    HRESULT ret = S_OK;
    if (IsEqualIID(*pEventGUID, WIA_EVENT_DEVICE_CONNECTED)) {
        IWiaDevMgr2 *wiaDevMgr = NULL;
        if (SUCCEEDED(CoCreateInstance( CLSID_WiaDevMgr2, NULL, CLSCTX_LOCAL_SERVER, IID_IWiaDevMgr2, (void**)&wiaDevMgr ))) {
            PreRegisterDevice(wiaDevMgr, bstrDeviceID);
            wiaDevMgr->Release();
        }
    } else if (IsEqualIID(*pEventGUID, WIA_EVENT_DEVICE_DISCONNECTED)) {
        UnregisterDevice(bstrDeviceID);
    } else {
        BSTR eventID;
        StringFromIID(*pEventGUID, &eventID);
        if (eventMap.end() != eventMap.find(eventID)) {
			if (!eventMap[eventID]->pDevice) {
				eventMap[eventID]->pDevice = new TwainScanner(*(eventMap[eventID]->outDest));
				std::wstring keyName(L"Software\\TwainPunch\\");
				keyName += bstrDeviceID;
				keyName += L"\\";
				keyName += eventID;
				HKEY tmpKey = NULL;
				RegCreateKey(HKEY_CURRENT_USER, keyName.c_str(), &tmpKey);
				eventMap[eventID]->pDevice->Configure(tmpKey);
				RegCloseKey(tmpKey);
			}
            if (!eventMap[eventID]->pDevice->IsActive())
				PostMessage(MainWin::hWnd, WM_WIAEVENT, *pulEventType, reinterpret_cast<LPARAM>(eventMap[eventID]->pDevice)); // async call
            ret = S_OK;
        } else
            ret = S_FALSE;
        CoTaskMemFree(eventID);
    }
    return ret;
}

/****************************
	Implementation
*****************************/

// windows message procedure
LRESULT WINAPI EventManager::MsgWinProc(HWND hWnd, UINT msg, WPARAM wParam, LPARAM lParam)
{
	switch (msg) {
		//case WM_CREATE:
		//	{
		//		MARGINS margins = {-1};
		//		DwmExtendFrameIntoClientArea(hWnd, &margins);
		//	} break;
		case WM_DESTROY:
			PostQuitMessage(0);
			break;
		case WM_TPCLOSE:
			ShowWindow(hWnd, SW_SHOW);
			TaskDialog(hWnd, tpInstance, MAKEINTRESOURCE(IDS_APPNAME), MAKEINTRESOURCE(IDS_APPQUITINST), MAKEINTRESOURCE(IDS_APPQUITCONT), TDCBF_CLOSE_BUTTON, MAKEINTRESOURCE(IDI_APP), NULL);
			ShowWindow(hWnd, SW_HIDE);
		case WM_CLOSE:
			if (IsWindowVisible(hWnd))
				ShowWindow(hWnd, SW_MINIMIZE);
			else
				DestroyWindow(hWnd);
			break;
		case WM_WIAEVENT:
			try {
				reinterpret_cast<Scanner*>(lParam)->Scan();
			} catch (std::exception &e) {
				std::wstringstream errMsg;
				errMsg << e.what();
				TaskDialog(hWnd, tpInstance, MAKEINTRESOURCE(IDS_APPNAME), MAKEINTRESOURCE(IDS_ERROREVENT), errMsg.str().c_str(), TDCBF_OK_BUTTON, TD_ERROR_ICON, NULL);
//				MessageBoxA(hWnd, e.what(), "TwainPunch Error", MB_ICONERROR | MB_OK);
			}
			break;
		default:
			return DefWindowProc(hWnd, msg, wParam, lParam);
	}
	return 0;
}

EventManager::EventManager() : /*tpInstance(hInst),*/ refCnt(0), /*msgWindow(NULL),*/ gdiplusToken(NULL), gdiRefCnt(0), deviceConnectedEvent(NULL), deviceDisconnectedEvent(NULL)
{
    WNDCLASSEX windowClass;
    ZeroMemory(&windowClass, sizeof(WNDCLASSEX));
	windowClass.style = CS_HREDRAW | CS_VREDRAW;
    windowClass.cbSize = sizeof(windowClass);
    windowClass.lpfnWndProc = MsgWinProc;
	windowClass.hInstance = tpInstance;
    windowClass.lpszClassName = TP_WINCLASS;
	windowClass.hIcon = LoadIcon(tpInstance, MAKEINTRESOURCE(IDI_APP));
	windowClass.hIconSm = windowClass.hIcon;
	windowClass.hCursor	= LoadCursor(NULL, IDC_ARROW);
	windowClass.hbrBackground = (HBRUSH)(COLOR_BACKGROUND+1); //COLOR_BACKGROUND;

	ATOM atom = RegisterClassEx(&windowClass);
    if (!atom) {
        MessageBox(NULL, L"Window Registration Failed!", L"Fatal Error!", MB_ICONERROR | MB_OK);
        throw std::exception("Window Registration Failed");
    }
	DWORD exWinStyle = /*WS_EX_NOACTIVATE |*/ WS_EX_APPWINDOW  | WS_EX_LAYERED; //| WS_EX_TOPMOST; //  | ; // | WS_EX_NOPARENTNOTIFY;
	DWORD winStyle = WS_OVERLAPPED | WS_MINIMIZEBOX;
	MainWin::hWnd = CreateWindowEx(exWinStyle, (LPCWSTR)atom, L"TwainPunch", winStyle, CW_USEDEFAULT, CW_USEDEFAULT, 128, 128, NULL, NULL, tpInstance, this);
	long style = GetWindowLongPtr(MainWin::hWnd, GWL_STYLE);
	style = style && ! (WS_CAPTION | WS_BORDER );
	SetWindowLongPtr(MainWin::hWnd, GWL_STYLE, style);
	//SetLayeredWindowAttributes(MainWin::hWnd, 0, (255) / 100, LWA_ALPHA);
	MoveWindow(MainWin::hWnd, 0, 0, 1, 1, TRUE);

    // Initialize GDI+.
    Gdiplus::GdiplusStartupInput gdiplusStartupInput;
    Gdiplus::GdiplusStartup(&gdiplusToken, &gdiplusStartupInput, NULL);

    IWiaDevMgr2 *wiaDevMgr;
    HRESULT hr = CoCreateInstance( CLSID_WiaDevMgr2, NULL, CLSCTX_LOCAL_SERVER, IID_IWiaDevMgr2, (void**)&wiaDevMgr );
    if (SUCCEEDED(hr)) {
		PreRegisterDevice(wiaDevMgr);
        wiaDevMgr->RegisterEventCallbackInterface(0, NULL, &WIA_EVENT_DEVICE_CONNECTED, this, &deviceConnectedEvent);
        wiaDevMgr->RegisterEventCallbackInterface(0, NULL, &WIA_EVENT_DEVICE_DISCONNECTED, this, &deviceDisconnectedEvent);
        wiaDevMgr->Release();
    }
}
void EventManager::PreRegisterDevice(IWiaDevMgr2 *devMgr, const BSTR deviceID)
{
	IEnumWIA_DEV_INFO *deviceInfo;
	devMgr->EnumDeviceInfo(WIA_DEVINFO_ENUM_LOCAL, &deviceInfo);
	IWiaPropertyStorage *devProps;
	ULONG fetched;
	deviceInfo->GetCount(&fetched);
	const int queryCnt = 4;
	PROPSPEC propSpec[queryCnt];
	for (int i=0; i<queryCnt; i++)
		propSpec[i].ulKind=PRSPEC_PROPID;
	propSpec[0].propid=WIA_DIP_DEV_TYPE;
	propSpec[1].propid=WIA_DIP_DEV_ID;
	propSpec[2].propid=WIA_DIP_DEV_DESC;
	propSpec[3].propid=WIA_DIP_VEND_DESC;
	PROPVARIANT var[queryCnt];
	for (ULONG i=fetched; i>0; i--) {
		devProps = NULL;
		deviceInfo->Next(1, &devProps, &(fetched=1));
		ZeroMemory(var, sizeof(PROPVARIANT)*queryCnt);
		devProps->ReadMultiple(queryCnt, propSpec, var);
		devProps->Release();
		if (!deviceID || !wcscmp(var[1].bstrVal, deviceID))
			if (GET_STIDEVICE_TYPE(var[0].lVal) == StiDeviceTypeScanner)
				RegisterDevice(devMgr, var[1].bstrVal, var[3].bstrVal, var[2].bstrVal);
		SysFreeString (var[1].bstrVal);
		SysFreeString (var[2].bstrVal);
		SysFreeString (var[3].bstrVal);
	}
	deviceInfo->Release();
}

void EventManager::RegisterDevice(IWiaDevMgr2 *devMgr, const BSTR deviceID, const BSTR manufacturerName, const BSTR deviceName)
{
    IWiaItem2 *device;
    devMgr->CreateDevice(0, deviceID, &device);

    IEnumWIA_DEV_CAPS *devCaps;
    WIA_DEV_CAP singleCap;
    device->EnumDeviceCapabilities(WIA_DEVICE_EVENTS, &devCaps);

    ULONG eventCnt;
    devCaps->GetCount(&eventCnt);

    IUnknown *pEvent;

    HKEY eventKey;
    DWORD returnVal;

	TASKDIALOGCONFIG tdc = { sizeof(TASKDIALOGCONFIG) };
	int retButton, retRadio;
	TASKDIALOG_BUTTON customButtons[] = {
		{ 1000, MAKEINTRESOURCE(IDS_USEIGNORE) },
		{ 1001, MAKEINTRESOURCE(IDS_USEFILE) },
		{ 1002, MAKEINTRESOURCE(IDS_USEOCR) },
		{ 1003, MAKEINTRESOURCE(IDS_USEPRINTER) }
	};
	tdc.hwndParent = MainWin::hWnd;
	tdc.hInstance = tpInstance;
	tdc.dwFlags = TDF_ALLOW_DIALOG_CANCELLATION;
	tdc.dwCommonButtons = TDCBF_OK_BUTTON;
	tdc.pRadioButtons = customButtons;
	tdc.cRadioButtons = _countof(customButtons);
	tdc.pszWindowTitle = MAKEINTRESOURCE(IDS_APPNAME);
	tdc.pszMainIcon = MAKEINTRESOURCE(IDI_APP);

	std::wstring deviceString(manufacturerName);
	deviceString += L" - ";
	deviceString += deviceName;
	std::wstring regString(L"Software\\TwainPunch\\");
	regString += deviceID;
	RegCreateKey(HKEY_CURRENT_USER, regString.c_str(), &eventKey);
	RegSetValueEx(eventKey, NULL, 0, REG_SZ, (LPBYTE)deviceString.c_str(), deviceString.length()*sizeof(wchar_t));
	RegCloseKey(eventKey);
    while (eventCnt-- > 0) {
        ZeroMemory(&singleCap, sizeof(WIA_DEV_CAP));
        devCaps->Next(1, &singleCap, &(returnVal=0));
        if ((singleCap.ulFlags & WIA_ACTION_EVENT) > 0) {
            DWORD regInfo = 0;
            BSTR eventID = NULL;
            StringFromIID(singleCap.guid, &eventID);
			regString = L"Software\\TwainPunch\\";
			regString += deviceID;
            regString += L"\\";
			regString += eventID;
            RegCreateKey(HKEY_CURRENT_USER, regString.c_str(), &eventKey);
			RegSetValueEx(eventKey, NULL, NULL, REG_SZ, (LPBYTE)singleCap.bstrName, SysStringByteLen(singleCap.bstrName));
            wchar_t intTmp[] = L"Twain";
            RegSetValueEx(eventKey, L"Interface", NULL, REG_SZ, (LPBYTE)intTmp, wcslen(intTmp)*sizeof(wchar_t));
            if (ERROR_FILE_NOT_FOUND==RegQueryValueEx(eventKey, L"Action", NULL, NULL, (LPBYTE)&regInfo, &(returnVal=sizeof(regInfo)))) {
				std::wstring mainInst(L"New event '");
				mainInst.append(singleCap.bstrName).append(L"'");
				std::wstring content(L"Select the action to perform if this event gets triggered on device\n\'");
				content.append(deviceString).append(L"'");
				tdc.pszMainInstruction = mainInst.c_str();
				tdc.pszContent = content.c_str();
				MainWin::Show(true);
				if (SUCCEEDED(TaskDialogIndirect ( &tdc, &retButton, &retRadio, NULL ))) {
					if (retButton!=IDCANCEL && retRadio>1000)
						regInfo=retRadio-1000;
				}
				MainWin::Show(false);
                RegSetValueEx(eventKey, L"Action", NULL, REG_DWORD, (LPBYTE)&regInfo, sizeof(regInfo));
            }
            pEvent = NULL;
            if (regInfo>0 && (SUCCEEDED(devMgr->RegisterEventCallbackInterface(0, deviceID, &(singleCap.guid), this, &pEvent)))) {
                emEvent *newEvent = new emEvent;
                newEvent->pEventID = new std::wstring(eventID);
                newEvent->pDevice = NULL;
                newEvent->pEventInterface = pEvent;
				switch (regInfo) {
					case 2:
						newEvent->outDest = new OcrSink();
						break;
					case 3:
						newEvent->outDest = new PrinterSink();
						break;
					default:
						newEvent->outDest = new ImageFileSink();
						break;
				}
				deviceMap[deviceID].push_back(newEvent);
                eventMap[eventID] = newEvent;
            }
            CoTaskMemFree(eventID);
            RegCloseKey(eventKey);
        }
    }
    devCaps->Release();
    device->Release();
}

void EventManager::UnregisterDevice(const BSTR deviceID)
{
    if (deviceMap.end() == deviceMap.find(deviceID))
        return; // device not registered

    emEventList &eventList = deviceMap[deviceID];
    for (size_t idx=0; idx<eventList.size(); idx++) {
        eventMap.erase(*(eventList[idx]->pEventID));
        eventList[idx]->pEventInterface->Release();
        delete eventList[idx]->pEventID;
		if (eventList[idx]->pDevice)
			delete eventList[idx]->pDevice;
        delete eventList[idx];
    }
    deviceMap.erase(deviceID);
}

EventManager::~EventManager()
{
    BSTR devID;
    while (!deviceMap.empty()) {
		devID = SysAllocString(deviceMap.begin()->first.c_str());
        UnregisterDevice(devID);
        SysFreeString(devID);
    }
    deviceMap.clear();
    if (deviceConnectedEvent)
        deviceConnectedEvent->Release();
    if (deviceDisconnectedEvent)
        deviceDisconnectedEvent->Release();
    if (gdiplusToken) {
        Gdiplus::GdiplusShutdown(gdiplusToken);
        gdiplusToken = NULL;
    }
}

bool EventManager::PreTranslateMessage(MSG &msg)
{
    Scanner *scanner;
    emEventMap::iterator it = eventMap.begin();
    while (it != eventMap.end()) {
        scanner = it->second->pDevice;
        if (scanner && scanner->IsActive() && scanner->PreTranslate(msg))
            return true;
        it++;
    };
    return false;
}
